package com.shangxia.basic.common.mybatis.dynamicsource;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.MutablePropertyValues;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.GenericBeanDefinition;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.env.Environment;
import org.springframework.core.type.AnnotationMetadata;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

/**
 * @ProjectName: MyDemo_zy
 * @ClassName: DynamicDataSourceRegister
 * @Time: 2020-04-01 20:33
 * @Author: zhou24
 * @Description: 初始化数据源
 * @Experience:
 */
@Slf4j
public class DynamicDataSourceRegister implements ImportBeanDefinitionRegistrar, EnvironmentAware {

    private static final String DATASOURCE_TYPE_DEFAULT = "org.apache.tomcat.jdbc.pool.DataSource";

    private DataSource defaultDataSource;

    private Map<String, DataSource> slaveDataSource = new HashMap<>();

    @Override
    public void setEnvironment(Environment environment) {
        initDefaultDataSource(environment);
        initSlaveDataSource(environment);
    }

    private void initDefaultDataSource(Environment env) {
        defaultDataSource = buildDataSource(env, "spring.datasource", ResourceTypeEnum.getResourceType(ResourceTypeEnum.BASE));
    }

    private void initSlaveDataSource(Environment env) {
        String dsPrefixs = env.getProperty("spring.datasource.names");
        for (String dsPrefix : dsPrefixs.split(",")) {
            ResourceType resourceType = ResourceTypeEnum.getResourceType(dsPrefix.toUpperCase());
            DataSource dataSource = buildDataSource(env, "spring.datasource." + dsPrefix, resourceType);
            slaveDataSource.put(resourceType.getKey(), dataSource);
        }
    }

    @Override
    public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata, BeanDefinitionRegistry registry) {
        Map<Object, Object> targetDataSources = new HashMap<>(16);
        targetDataSources.put(ResourceTypeEnum.BASE, this.defaultDataSource);
        DynamicDataSourceContextHolder.dataSourceIds.add(ResourceTypeEnum.BASE);
        targetDataSources.putAll(slaveDataSource);
        slaveDataSource.keySet().forEach(key -> DynamicDataSourceContextHolder.dataSourceIds.add(key));
        GenericBeanDefinition genericBeanDefinition = new GenericBeanDefinition();
        genericBeanDefinition.setBeanClass(DynamicDataSource.class);
        genericBeanDefinition.setSynthetic(true);
        MutablePropertyValues mpv = genericBeanDefinition.getPropertyValues();
        mpv.addPropertyValue("defaultTargetDataSource", defaultDataSource);
        mpv.addPropertyValue("targetDataSources", targetDataSources);
        registry.registerBeanDefinition("dataSource", genericBeanDefinition);
    }

    private DataSource buildDataSource(Environment env, String prefix, ResourceType rte) {
        Map<String, String> dataSourceMap = new HashMap<>(16);
        if (env.containsProperty(prefix + ".driver-class-name")) {
            dataSourceMap.put("driver", env.getProperty(prefix + ".driver-class-name"));
        }
        dataSourceMap.put("url", env.getProperty(prefix + ".url"));
        dataSourceMap.put("username", env.getProperty(prefix + ".username"));
        dataSourceMap.put("password", env.getProperty(prefix + ".password"));

        try {
            String type = dataSourceMap.get("type");
            if (null == type) {
                type = DATASOURCE_TYPE_DEFAULT;
            }
            Class<? extends DataSource> dataSourceType;
            dataSourceType = (Class<? extends DataSource>) Class.forName(type);
            String driverClassName = dataSourceMap.get("driver");
            String url = dataSourceMap.get("url");
            String username = dataSourceMap.get("username");
            String password = dataSourceMap.get("password");
            DataSource dataSource = DataSourceBuilder.create().driverClassName(driverClassName).url(url).username(username).password(password).type(dataSourceType).build();
            ((org.apache.tomcat.jdbc.pool.DataSource) dataSource).setTestWhileIdle(true);
            ((org.apache.tomcat.jdbc.pool.DataSource) dataSource).setMaxIdle(180000);
            ((org.apache.tomcat.jdbc.pool.DataSource) dataSource).setMaxActive(10);
            ((org.apache.tomcat.jdbc.pool.DataSource) dataSource).setInitialSize(2);
            ((org.apache.tomcat.jdbc.pool.DataSource) dataSource).setMaxIdle(180000);
            ((org.apache.tomcat.jdbc.pool.DataSource) dataSource).setValidationInterval(120000);
            switch (driverClassName) {
                case "oracle.jdbc.OracleDriver":
                    ((org.apache.tomcat.jdbc.pool.DataSource) dataSource).setValidationQuery("select 1 from dual");
                    break;
                default:
                    ((org.apache.tomcat.jdbc.pool.DataSource) dataSource).setValidationQuery("select 1");
            }
            rte.setDriver(driverClassName);
            return dataSource;
        } catch (ClassNotFoundException e) {
            log.error("e:{}", e);
        }
        return null;
    }


}
